use std::env;
use std::mem;
use std::path::Path;
use windows::Win32::Foundation::HANDLE;
use windows::Win32::System::Diagnostics::ToolHelp::{
    CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW, TH32CS_SNAPPROCESS,
};
use windows::Win32::System::Threading::{OpenProcess, PROCESS_QUERY_INFORMATION, PROCESS_VM_READ};

mod classes;
mod enums;
mod get_process_mitigation_command;
mod ntdll;
mod registry_manager;
mod set_process_mitigation_command;

use classes::*;
use enums::*;
use get_process_mitigation_command::*;
use registry_manager::*;
use set_process_mitigation_command::*;

#[derive(Debug, Clone)]
pub struct ProcessInfo {
    pub process_name: String,
    pub id: u32,
    pub handle: Option<HANDLE>,
}

impl ProcessInfo {
    pub fn get_handle(&mut self) -> Result<HANDLE, windows::core::Error> {
        if let Some(handle) = self.handle {
            Ok(handle)
        } else {
            unsafe {
                let handle =
                    OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, false, self.id)?;
                self.handle = Some(handle);
                Ok(handle)
            }
        }
    }
}

/// Represents the result of an operation that can either succeed or fail.
#[derive(Debug, Clone)]
pub struct AppResult<T> {
    pub is_success: bool,
    pub value: Option<T>,
    pub error: String,
}

impl<T> AppResult<T> {
    pub fn success(value: T) -> Self {
        Self {
            is_success: true,
            value: Some(value),
            error: String::new(),
        }
    }

    pub fn failure(error: &str) -> Self {
        Self {
            is_success: false,
            value: None,
            error: error.to_string(),
        }
    }

    pub fn is_failure(&self) -> bool {
        !self.is_success
    }

    pub fn match_result<R>(
        &self,
        on_success: impl FnOnce(&T) -> R,
        on_failure: impl FnOnce(&str) -> R,
    ) -> R {
        if self.is_success {
            if let Some(ref value) = self.value {
                on_success(value)
            } else {
                on_failure("Success result has no value")
            }
        } else {
            on_failure(&self.error)
        }
    }

    pub fn get_value_or_default(&self, default_value: T) -> T
    where
        T: Clone,
    {
        if self.is_success {
            self.value.clone().unwrap_or(default_value)
        } else {
            default_value
        }
    }
}

/// Result type for operations that don't return a value
#[derive(Debug, Clone)]
pub struct AppResultVoid {
    pub is_success: bool,
    pub error: String,
}

impl AppResultVoid {
    pub fn success() -> Self {
        Self {
            is_success: true,
            error: String::new(),
        }
    }

    pub fn failure(error: &str) -> Self {
        Self {
            is_success: false,
            error: error.to_string(),
        }
    }

    pub fn is_failure(&self) -> bool {
        !self.is_success
    }

    pub fn match_result<R>(
        &self,
        on_success: impl FnOnce() -> R,
        on_failure: impl FnOnce(&str) -> R,
    ) -> R {
        if self.is_success {
            on_success()
        } else {
            on_failure(&self.error)
        }
    }
}

fn parse_mitigation_numbers(numbers_str: &str) -> Result<Vec<MitigationOptions>, String> {
    if numbers_str.trim().is_empty() {
        return Ok(Vec::new());
    }

    let mut options = Vec::new();
    for num_str in numbers_str.split(',') {
        let num_str = num_str.trim();
        if !num_str.is_empty() {
            match num_str.parse::<u32>() {
                Ok(num) => match MitigationOptions::from_number(num) {
                    Ok(option) => options.push(option),
                    Err(e) => return Err(e),
                },
                Err(_) => return Err(format!("Invalid number format: {}", num_str)),
            }
        }
    }
    Ok(options)
}

fn parse_string_list(strings_str: &str) -> Vec<String> {
    if strings_str.trim().is_empty() {
        return Vec::new();
    }

    strings_str
        .split(',')
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
        .collect()
}

pub struct Program;

impl Program {
    /// Gets the current settings on all running instances of a specific process.
    pub fn get_from_running_process_name(
        running_process_name: &str,
    ) -> AppResult<Vec<ProcessMitigations>> {
        let process_name = Path::new(running_process_name)
            .file_stem()
            .and_then(|s| s.to_str())
            .unwrap_or(running_process_name);

        let processes = Self::get_processes_by_name(process_name);
        let mut output = Vec::new();

        for process in processes {
            match GetProcessMitigationCommand::get_policy_from_running_process(&process) {
                Ok(mitigation) => output.push(mitigation),
                Err(e) => {
                    return AppResult::failure(&format!(
                        "Failed to get policies from running process: {}",
                        e
                    ));
                }
            }
        }

        AppResult::success(output)
    }

    /// Gets the current system process mitigation defaults stored in the registry.
    pub fn get_system_policy() -> AppResult<AppMitigations> {
        match std::panic::catch_unwind(|| {
            let registry_manager = RegistryManager::new();
            registry_manager.get_system_policy()
        }) {
            Ok(Ok(policy)) => AppResult::success(policy),
            Ok(Err(error)) => AppResult::failure(&error),
            Err(_) => AppResult::failure("Failed to get system policy"),
        }
    }

    /// Gets all policies for all processes set in the registry.
    pub fn get_full_policies_for_all_processes() -> AppResult<Vec<AppMitigations>> {
        match std::panic::catch_unwind(|| {
            let registry_manager = RegistryManager::new();
            registry_manager.get_policy_list_from_registry(None)
        }) {
            Ok(Ok(policies)) => AppResult::success(policies),
            Ok(Err(error)) => AppResult::failure(&error),
            Err(_) => AppResult::failure("Failed to get full policies for all processes"),
        }
    }

    /// Gets the current settings in the registry for a specific process
    pub fn get_policy_from_registry_by_name(process_name: &str) -> AppResult<Vec<AppMitigations>> {
        match std::panic::catch_unwind(|| {
            let registry_manager = RegistryManager::new();
            registry_manager.get_policy_from_registry_by_name(process_name)
        }) {
            Ok(Ok(policies)) => AppResult::success(policies),
            Ok(Err(error)) => AppResult::failure(&error),
            Err(_) => AppResult::failure("Failed to get policy from registry by name"),
        }
    }

    /// Add mitigations for the Process
    pub fn add_mitigations_for_process(
        process_name: &str,
        disable_list: Option<&[MitigationOptions]>,
        enable_list: Option<&[MitigationOptions]>,
        eaf_modules_list: Option<&[String]>,
        is_force: Option<&str>,
        is_remove: bool,
        is_reset: bool,
    ) -> AppResultVoid {
        let mut registry_manager = RegistryManager::new();

        if enable_list.is_none() && disable_list.is_none() {
            return AppResultVoid::failure(
                "The input parameters are invalid. Please specify a list of mitigations to enable or disable for this process.",
            );
        }

        let mut process_name = process_name.to_string();
        if process_name
            == Path::new(&process_name)
                .file_stem()
                .and_then(|s| s.to_str())
                .unwrap_or(&process_name)
        {
            process_name = format!("{}.exe", process_name).to_lowercase();
        }

        let from_registry_by_name =
            match registry_manager.get_policy_from_registry_by_name(&process_name) {
                Ok(policies) => policies,
                Err(error) => return AppResultVoid::failure(&error),
            };

        let mut app_mitigations = if from_registry_by_name.is_empty() {
            AppMitigations::new(&process_name)
        } else if from_registry_by_name.len() > 1 {
            return AppResultVoid::failure(
                "Multiple mitigation policies found that may match the given process name. Please specify the full path to be matched instead.",
            );
        } else {
            from_registry_by_name.into_iter().next().unwrap()
        };

        match SetProcessMitigationsCommand::set_enable_and_disable(
            &mut app_mitigations,
            &mut registry_manager,
            disable_list,
            enable_list,
            eaf_modules_list,
            is_force,
            is_remove,
            is_reset,
            false, // isSystemMode
        ) {
            Ok(_) => AppResultVoid::success(),
            Err(e) => {
                AppResultVoid::failure(&format!("Failed to add mitigations for process: {}", e))
            }
        }
    }

    /// Add mitigations to the System
    pub fn add_mitigations_to_system(
        disable_list: Option<&[MitigationOptions]>,
        enable_list: Option<&[MitigationOptions]>,
        eaf_modules_list: Option<&[String]>,
        is_force: Option<&str>,
        is_remove: bool,
        is_reset: bool,
    ) -> AppResultVoid {
        let mut registry_manager = RegistryManager::new();

        let mut app_mitigations = match registry_manager.get_system_policy() {
            Ok(policy) => policy,
            Err(_) => {
                let mut policy = AppMitigations::new("System");
                policy.source = "System Defaults".to_string();
                policy
            }
        };

        match SetProcessMitigationsCommand::set_enable_and_disable(
            &mut app_mitigations,
            &mut registry_manager,
            disable_list,
            enable_list,
            eaf_modules_list,
            is_force,
            is_remove,
            is_reset,
            true, // isSystemMode
        ) {
            Ok(_) => AppResultVoid::success(),
            Err(e) => {
                AppResultVoid::failure(&format!("Failed to add mitigations to system: {}", e))
            }
        }
    }

    /// Import mitigations from XML file
    pub fn import_xml_mitigations(file_path: &str) -> AppResultVoid {
        let result = SetProcessMitigationsCommand::import_mitigation(file_path);
        if result >= 0 {
            AppResultVoid::success()
        } else {
            AppResultVoid::failure(&format!("Error importing: {}", result))
        }
    }

    /// Validate Mitigations XML file
    pub fn validate_xml_mitigations(file_path: &str) -> AppResult<bool> {
        let mut result = false;
        let num = SetProcessMitigationsCommand::validate_xml_from_managed(file_path, &mut result);
        if num < 0 {
            AppResult::failure(&format!("Error validating: {}", num))
        } else {
            println!("Is Valid XML: {}", result);
            AppResult::success(result)
        }
    }

    /// Add mitigations for the Process using comma-separated numbers
    pub fn add_mitigations_for_process_ex(
        process_name: &str,
        disable_list_numbers: &str,
        enable_list_numbers: &str,
        eaf_modules_list: &str,
        is_force: Option<&str>,
        is_remove: bool,
        is_reset: bool,
    ) -> AppResultVoid {
        let disable_options = match parse_mitigation_numbers(disable_list_numbers) {
            Ok(options) => {
                if options.is_empty() {
                    None
                } else {
                    Some(options)
                }
            }
            Err(e) => return AppResultVoid::failure(&format!("Error parsing disable list: {}", e)),
        };

        let enable_options = match parse_mitigation_numbers(enable_list_numbers) {
            Ok(options) => {
                if options.is_empty() {
                    None
                } else {
                    Some(options)
                }
            }
            Err(e) => return AppResultVoid::failure(&format!("Error parsing enable list: {}", e)),
        };

        let eaf_modules = parse_string_list(eaf_modules_list);
        let eaf_modules_ref = if eaf_modules.is_empty() {
            None
        } else {
            Some(eaf_modules.as_slice())
        };

        Self::add_mitigations_for_process(
            process_name,
            disable_options.as_deref(),
            enable_options.as_deref(),
            eaf_modules_ref,
            is_force,
            is_remove,
            is_reset,
        )
    }

    /// Add mitigations to the System using comma-separated numbers
    pub fn add_mitigations_to_system_ex(
        disable_list_numbers: &str,
        enable_list_numbers: &str,
        eaf_modules_list: &str,
        is_force: Option<&str>,
        is_remove: bool,
        is_reset: bool,
    ) -> AppResultVoid {
        let disable_options = match parse_mitigation_numbers(disable_list_numbers) {
            Ok(options) => {
                if options.is_empty() {
                    None
                } else {
                    Some(options)
                }
            }
            Err(e) => return AppResultVoid::failure(&format!("Error parsing disable list: {}", e)),
        };

        let enable_options = match parse_mitigation_numbers(enable_list_numbers) {
            Ok(options) => {
                if options.is_empty() {
                    None
                } else {
                    Some(options)
                }
            }
            Err(e) => return AppResultVoid::failure(&format!("Error parsing enable list: {}", e)),
        };

        let eaf_modules = parse_string_list(eaf_modules_list);
        let eaf_modules_ref = if eaf_modules.is_empty() {
            None
        } else {
            Some(eaf_modules.as_slice())
        };

        Self::add_mitigations_to_system(
            disable_options.as_deref(),
            enable_options.as_deref(),
            eaf_modules_ref,
            is_force,
            is_remove,
            is_reset,
        )
    }

    /// Gets the current settings in the registry for a specific process and returns JSON
    pub fn get_policy_from_registry_by_name_ex(process_name: &str) -> AppResult<String> {
        match Self::get_policy_from_registry_by_name(process_name) {
            app_result if app_result.is_success => {
                if let Some(policies) = app_result.value {
                    let json_policies: Vec<AppMitigationsJson> =
                        policies.iter().map(|p| p.to_json()).collect();
                    match serde_json::to_string_pretty(&json_policies) {
                        Ok(json) => AppResult::success(json),
                        Err(e) => {
                            AppResult::failure(&format!("Failed to serialize to JSON: {}", e))
                        }
                    }
                } else {
                    AppResult::success("[]".to_string())
                }
            }
            app_result => AppResult::failure(&app_result.error),
        }
    }

    /// Gets the current system process mitigation defaults stored in the registry and returns JSON
    pub fn get_system_policy_ex() -> AppResult<String> {
        match Self::get_system_policy() {
            app_result if app_result.is_success => {
                if let Some(policy) = app_result.value {
                    let json_policy = policy.to_json();
                    match serde_json::to_string_pretty(&json_policy) {
                        Ok(json) => AppResult::success(json),
                        Err(e) => {
                            AppResult::failure(&format!("Failed to serialize to JSON: {}", e))
                        }
                    }
                } else {
                    AppResult::failure("No system policy found")
                }
            }
            app_result => AppResult::failure(&app_result.error),
        }
    }

    fn get_processes_by_name(process_name: &str) -> Vec<ProcessInfo> {
        let mut processes = Vec::new();

        unsafe {
            let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
            if snapshot.is_err() {
                return processes;
            }
            let snapshot = snapshot.unwrap();

            let mut process_entry = PROCESSENTRY32W {
                dwSize: mem::size_of::<PROCESSENTRY32W>() as u32,
                ..Default::default()
            };

            if Process32FirstW(snapshot, &mut process_entry).is_ok() {
                loop {
                    let exe_file = String::from_utf16_lossy(&process_entry.szExeFile)
                        .trim_end_matches('\0')
                        .to_string();

                    if exe_file.eq_ignore_ascii_case(&format!("{}.exe", process_name))
                        || exe_file.eq_ignore_ascii_case(process_name)
                    {
                        processes.push(ProcessInfo {
                            process_name: exe_file,
                            id: process_entry.th32ProcessID,
                            handle: None,
                        });
                    }

                    if Process32NextW(snapshot, &mut process_entry).is_err() {
                        break;
                    }
                }
            }
        }

        processes
    }
}

// Convert OptionValue to readable string
fn option_value_to_string(value: OptionValue) -> &'static str {
    match value {
        OptionValue::NotSet => "Not Set",
        OptionValue::On => "On",
        OptionValue::Off => "Off",
        OptionValue::Error => "Error",
    }
}

// to print ProcessMitigations details
fn print_process_mitigations(mitigations: &ProcessMitigations) {
    println!("========================================");
    println!("Process: {}", mitigations.process_name);
    println!("Source: {}", mitigations.source);
    println!("ID: {}", mitigations.id);
    println!("========================================");

    // DEP Policy
    if let Some(dep) = mitigations.dep() {
        println!("\n--- DEP Policy ---");
        println!("  Enable: {}", option_value_to_string(dep.enable()));
        println!(
            "  Emulate ATL Thunks: {}",
            option_value_to_string(dep.emulate_atl_thunks())
        );
        println!("  Flags: {}", dep.flags());
    }

    // ASLR Policy
    if let Some(aslr) = mitigations.aslr() {
        println!("\n--- ASLR Policy ---");
        println!(
            "  Force Relocate Images: {}",
            option_value_to_string(aslr.force_relocate_images())
        );
        println!(
            "  Require Info: {}",
            option_value_to_string(aslr.require_info())
        );
        println!("  Bottom Up: {}", option_value_to_string(aslr.bottom_up()));
        println!(
            "  High Entropy: {}",
            option_value_to_string(aslr.high_entropy())
        );
        println!("  Flags: {}", aslr.flags());
    }

    // Strict Handle Policy
    if let Some(strict_handle) = mitigations.strict_handle() {
        println!("\n--- Strict Handle Policy ---");
        println!(
            "  Enable: {}",
            option_value_to_string(strict_handle.enable())
        );
        println!("  Flags: {}", strict_handle.flags());
    }

    // System Call Policy
    if let Some(system_call) = mitigations.system_call() {
        println!("\n--- System Call Policy ---");
        println!(
            "  Disable Win32k System Calls: {}",
            option_value_to_string(system_call.disable_win32k_system_calls())
        );
        println!("  Audit: {}", option_value_to_string(system_call.audit()));
        println!(
            "  Disable Fsctl System Calls: {}",
            option_value_to_string(system_call.disable_fsctl_system_calls())
        );
        println!(
            "  Audit Fsctl System Calls: {}",
            option_value_to_string(system_call.audit_fsctl_system_calls())
        );
        println!("  Flags: {}", system_call.flags());
    }

    // Extension Point Policy
    if let Some(extension_point) = mitigations.extension_point() {
        println!("\n--- Extension Point Policy ---");
        println!(
            "  Disable Extension Points: {}",
            option_value_to_string(extension_point.disable_extension_points())
        );
        println!("  Flags: {}", extension_point.flags());
    }

    // Dynamic Code Policy
    if let Some(dynamic_code) = mitigations.dynamic_code() {
        println!("\n--- Dynamic Code Policy ---");
        println!(
            "  Block Dynamic Code: {}",
            option_value_to_string(dynamic_code.block_dynamic_code())
        );
        println!(
            "  Allow Threads to Opt Out: {}",
            option_value_to_string(dynamic_code.allow_threads_to_opt_out())
        );
        println!("  Audit: {}", option_value_to_string(dynamic_code.audit()));
        println!("  Flags: {}", dynamic_code.flags());
    }

    // CFG Policy
    if let Some(cfg) = mitigations.cfg() {
        println!("\n--- Control Flow Guard Policy ---");
        println!("  Enable: {}", option_value_to_string(cfg.enable()));
        println!(
            "  Suppress Exports: {}",
            option_value_to_string(cfg.suppress_exports())
        );
        println!(
            "  Strict Control Flow Guard: {}",
            option_value_to_string(cfg.strict_control_flow_guard())
        );
        println!("  Flags: {}", cfg.flags());
    }

    // Binary Signature Policy
    if let Some(binary_signature) = mitigations.binary_signature() {
        println!("\n--- Binary Signature Policy ---");
        println!(
            "  Microsoft Signed Only: {}",
            option_value_to_string(binary_signature.microsoft_signed_only())
        );
        println!(
            "  Allow Store Signed Binaries: {}",
            option_value_to_string(binary_signature.allow_store_signed_binaries())
        );
        println!(
            "  Audit: {}",
            option_value_to_string(binary_signature.audit())
        );
        println!(
            "  Audit Store Signed: {}",
            option_value_to_string(binary_signature.audit_store_signed())
        );
        println!("  Flags: {}", binary_signature.flags());
    }

    // Font Disable Policy
    if let Some(font_disable) = mitigations.font_disable() {
        println!("\n--- Font Disable Policy ---");
        println!(
            "  Disable Non System Fonts: {}",
            option_value_to_string(font_disable.disable_non_system_fonts())
        );
        println!("  Audit: {}", option_value_to_string(font_disable.audit()));
        println!("  Flags: {}", font_disable.flags());
    }

    // Image Load Policy
    if let Some(image_load) = mitigations.image_load() {
        println!("\n--- Image Load Policy ---");
        println!(
            "  Block Remote Image Loads: {}",
            option_value_to_string(image_load.block_remote_image_loads())
        );
        println!(
            "  Block Low Label Image Loads: {}",
            option_value_to_string(image_load.block_low_label_image_loads())
        );
        println!(
            "  Prefer System32: {}",
            option_value_to_string(image_load.prefer_system32())
        );
        println!(
            "  Audit Remote Image Loads: {}",
            option_value_to_string(image_load.audit_remote_image_loads())
        );
        println!(
            "  Audit Low Label Image Loads: {}",
            option_value_to_string(image_load.audit_low_label_image_loads())
        );
        println!(
            "  Audit Prefer System32: {}",
            option_value_to_string(image_load.audit_prefer_system32())
        );
        println!("  Flags: {}", image_load.flags());
    }

    // Payload Policy
    if let Some(payload) = mitigations.payload() {
        println!("\n--- Payload Policy ---");
        println!(
            "  Enable Export Address Filter: {}",
            option_value_to_string(payload.enable_export_address_filter())
        );
        println!(
            "  Audit Enable Export Address Filter: {}",
            option_value_to_string(payload.audit_enable_export_address_filter())
        );
        println!(
            "  Enable Export Address Filter Plus: {}",
            option_value_to_string(payload.enable_export_address_filter_plus())
        );
        println!(
            "  Audit Enable Export Address Filter Plus: {}",
            option_value_to_string(payload.audit_enable_export_address_filter_plus())
        );
        println!(
            "  Enable Import Address Filter: {}",
            option_value_to_string(payload.enable_import_address_filter())
        );
        println!(
            "  Audit Enable Import Address Filter: {}",
            option_value_to_string(payload.audit_enable_import_address_filter())
        );
        println!(
            "  Enable ROP Stack Pivot: {}",
            option_value_to_string(payload.enable_rop_stack_pivot())
        );
        println!(
            "  Audit Enable ROP Stack Pivot: {}",
            option_value_to_string(payload.audit_enable_rop_stack_pivot())
        );
        println!(
            "  Enable ROP Caller Check: {}",
            option_value_to_string(payload.enable_rop_caller_check())
        );
        println!(
            "  Audit Enable ROP Caller Check: {}",
            option_value_to_string(payload.audit_enable_rop_caller_check())
        );
        println!(
            "  Enable ROP Sim Exec: {}",
            option_value_to_string(payload.enable_rop_sim_exec())
        );
        println!(
            "  Audit Enable ROP Sim Exec: {}",
            option_value_to_string(payload.audit_enable_rop_sim_exec())
        );
        println!("  Flags: {}", payload.flags());
    }

    // Child Process Policy
    if let Some(child_process) = mitigations.child_process() {
        println!("\n--- Child Process Policy ---");
        println!(
            "  Disallow Child Process Creation: {}",
            option_value_to_string(child_process.disallow_child_process_creation())
        );
        println!("  Audit: {}", option_value_to_string(child_process.audit()));
        println!("  Flags: {}", child_process.flags());
    }

    // User Shadow Stack Policy
    if let Some(user_shadow_stack) = mitigations.user_shadow_stack() {
        println!("\n--- User Shadow Stack Policy ---");
        println!(
            "  User Shadow Stack: {}",
            option_value_to_string(user_shadow_stack.user_shadow_stack())
        );
        println!(
            "  User Shadow Stack Strict Mode: {}",
            option_value_to_string(user_shadow_stack.user_shadow_stack_strict_mode())
        );
        println!(
            "  Audit User Shadow Stack: {}",
            option_value_to_string(user_shadow_stack.audit_user_shadow_stack())
        );
        println!(
            "  Set Context IP Validation: {}",
            option_value_to_string(user_shadow_stack.set_context_ip_validation())
        );
        println!(
            "  Audit Set Context IP Validation: {}",
            option_value_to_string(user_shadow_stack.audit_set_context_ip_validation())
        );
        println!(
            "  Block Non CET Binaries: {}",
            option_value_to_string(user_shadow_stack.block_non_cet_binaries())
        );
        println!(
            "  Block Non CET Binaries Non Ehcont: {}",
            option_value_to_string(user_shadow_stack.block_non_cet_binaries_non_ehcont())
        );
        println!(
            "  Audit Block Non CET Binaries: {}",
            option_value_to_string(user_shadow_stack.audit_block_non_cet_binaries())
        );
        println!("  Flags: {}", user_shadow_stack.flags());
    }

    // SEHOP Policy
    if let Some(sehop) = mitigations.sehop() {
        println!("\n--- SEHOP Policy ---");
        println!("  Enable: {}", option_value_to_string(sehop.enable()));
        println!(
            "  Telemetry Only: {}",
            option_value_to_string(sehop.telemetry_only())
        );
        println!("  Audit: {}", option_value_to_string(sehop.audit()));
        println!("  Flags: {}", sehop.flags());
    }

    println!("========================================\n");
}

// to print AppMitigations details
fn print_app_mitigations(mitigations: &AppMitigations) {
    println!("========================================");
    println!("Process: {}", mitigations.process_name);
    println!("Source: {}", mitigations.source);
    println!("========================================");

    // DEP Policy
    if let Some(dep) = mitigations.dep() {
        println!("\n--- DEP Policy ---");
        println!("  Enable: {}", option_value_to_string(dep.enable()));
        println!(
            "  Emulate ATL Thunks: {}",
            option_value_to_string(dep.emulate_atl_thunks())
        );
        println!("  Override DEP: {}", dep.override_dep());
        println!("  Flags: {}", dep.flags());
    }

    // ASLR Policy
    if let Some(aslr) = mitigations.aslr() {
        println!("\n--- ASLR Policy ---");
        println!(
            "  Force Relocate Images: {}",
            option_value_to_string(aslr.force_relocate_images())
        );
        println!(
            "  Require Info: {}",
            option_value_to_string(aslr.require_info())
        );
        println!("  Bottom Up: {}", option_value_to_string(aslr.bottom_up()));
        println!(
            "  High Entropy: {}",
            option_value_to_string(aslr.high_entropy())
        );
        println!(
            "  Override Force Relocate Images: {}",
            aslr.override_force_relocate_images()
        );
        println!("  Override Bottom Up: {}", aslr.override_bottom_up());
        println!("  Override High Entropy: {}", aslr.override_high_entropy());
        println!("  Flags: {}", aslr.flags());
    }

    // Strict Handle Policy
    if let Some(strict_handle) = mitigations.strict_handle() {
        println!("\n--- Strict Handle Policy ---");
        println!(
            "  Enable: {}",
            option_value_to_string(strict_handle.enable())
        );
        println!(
            "  Override Strict Handle: {}",
            strict_handle.override_strict_handle()
        );
        println!("  Flags: {}", strict_handle.flags());
    }

    // System Call Policy
    if let Some(system_call) = mitigations.system_call() {
        println!("\n--- System Call Policy ---");
        println!(
            "  Disable Win32k System Calls: {}",
            option_value_to_string(system_call.disable_win32k_system_calls())
        );
        println!("  Audit: {}", option_value_to_string(system_call.audit()));
        println!(
            "  Disable Fsctl System Calls: {}",
            option_value_to_string(system_call.disable_fsctl_system_calls())
        );
        println!(
            "  Audit Fsctl System Calls: {}",
            option_value_to_string(system_call.audit_fsctl_system_calls())
        );
        println!(
            "  Override System Call: {}",
            system_call.override_system_call()
        );
        println!(
            "  Override Fsctl System Call: {}",
            system_call.override_fsctl_system_call()
        );
        println!("  Flags: {}", system_call.flags());
    }

    // Extension Point Policy
    if let Some(extension_point) = mitigations.extension_point() {
        println!("\n--- Extension Point Policy ---");
        println!(
            "  Disable Extension Points: {}",
            option_value_to_string(extension_point.disable_extension_points())
        );
        println!(
            "  Override Extension Point: {}",
            extension_point.override_extension_point()
        );
        println!("  Flags: {}", extension_point.flags());
    }

    // Dynamic Code Policy
    if let Some(dynamic_code) = mitigations.dynamic_code() {
        println!("\n--- Dynamic Code Policy ---");
        println!(
            "  Block Dynamic Code: {}",
            option_value_to_string(dynamic_code.block_dynamic_code())
        );
        println!(
            "  Allow Threads to Opt Out: {}",
            option_value_to_string(dynamic_code.allow_threads_to_opt_out())
        );
        println!("  Audit: {}", option_value_to_string(dynamic_code.audit()));
        println!(
            "  Override Dynamic Code: {}",
            dynamic_code.override_dynamic_code()
        );
        println!("  Flags: {}", dynamic_code.flags());
    }

    // CFG Policy
    if let Some(cfg) = mitigations.cfg() {
        println!("\n--- Control Flow Guard Policy ---");
        println!("  Enable: {}", option_value_to_string(cfg.enable()));
        println!(
            "  Suppress Exports: {}",
            option_value_to_string(cfg.suppress_exports())
        );
        println!(
            "  Strict Control Flow Guard: {}",
            option_value_to_string(cfg.strict_control_flow_guard())
        );
        println!("  Override CFG: {}", cfg.override_cfg());
        println!("  Override Strict CFG: {}", cfg.override_strict_cfg());
        println!("  Flags: {}", cfg.flags());
    }

    // Binary Signature Policy
    if let Some(binary_signature) = mitigations.binary_signature() {
        println!("\n--- Binary Signature Policy ---");
        println!(
            "  Microsoft Signed Only: {}",
            option_value_to_string(binary_signature.microsoft_signed_only())
        );
        println!(
            "  Allow Store Signed Binaries: {}",
            option_value_to_string(binary_signature.allow_store_signed_binaries())
        );
        println!(
            "  Enforce Module Dependency Signing: {}",
            option_value_to_string(binary_signature.enforce_module_dependency_signing())
        );
        println!(
            "  Audit: {}",
            option_value_to_string(binary_signature.audit())
        );
        println!(
            "  Audit Store Signed: {}",
            option_value_to_string(binary_signature.audit_store_signed())
        );
        println!(
            "  Audit Enforce Module Dependency Signing: {}",
            option_value_to_string(binary_signature.audit_enforce_module_dependency_signing())
        );
        println!(
            "  Override Microsoft Signed Only: {}",
            binary_signature.override_microsoft_signed_only()
        );
        println!(
            "  Override Enforce Module Dependency Signing: {}",
            binary_signature.override_enforce_module_dependency_signing()
        );
        println!("  Flags: {}", binary_signature.flags());
    }

    // Font Disable Policy
    if let Some(font_disable) = mitigations.font_disable() {
        println!("\n--- Font Disable Policy ---");
        println!(
            "  Disable Non System Fonts: {}",
            option_value_to_string(font_disable.disable_non_system_fonts())
        );
        println!("  Audit: {}", option_value_to_string(font_disable.audit()));
        println!(
            "  Override Font Disable: {}",
            font_disable.override_font_disable()
        );
        println!("  Flags: {}", font_disable.flags());
    }

    // Image Load Policy
    if let Some(image_load) = mitigations.image_load() {
        println!("\n--- Image Load Policy ---");
        println!(
            "  Block Remote Image Loads: {}",
            option_value_to_string(image_load.block_remote_image_loads())
        );
        println!(
            "  Block Low Label Image Loads: {}",
            option_value_to_string(image_load.block_low_label_image_loads())
        );
        println!(
            "  Prefer System32: {}",
            option_value_to_string(image_load.prefer_system32())
        );
        println!(
            "  Audit Remote Image Loads: {}",
            option_value_to_string(image_load.audit_remote_image_loads())
        );
        println!(
            "  Audit Low Label Image Loads: {}",
            option_value_to_string(image_load.audit_low_label_image_loads())
        );
        println!(
            "  Audit Prefer System32: {}",
            option_value_to_string(image_load.audit_prefer_system32())
        );
        println!(
            "  Override Block Remote Image Loads: {}",
            image_load.override_block_remote_image_loads()
        );
        println!(
            "  Override Block Low Label: {}",
            image_load.override_block_low_label()
        );
        println!(
            "  Override Prefer System32: {}",
            image_load.override_prefer_system32()
        );
        println!("  Flags: {}", image_load.flags());
    }

    // Payload Policy
    if let Some(payload) = mitigations.payload() {
        println!("\n--- Payload Policy ---");
        println!(
            "  Enable Export Address Filter: {}",
            option_value_to_string(payload.enable_export_address_filter())
        );
        println!(
            "  Audit Enable Export Address Filter: {}",
            option_value_to_string(payload.audit_enable_export_address_filter())
        );
        println!(
            "  Enable Export Address Filter Plus: {}",
            option_value_to_string(payload.enable_export_address_filter_plus())
        );
        println!(
            "  Audit Enable Export Address Filter Plus: {}",
            option_value_to_string(payload.audit_enable_export_address_filter_plus())
        );
        println!(
            "  Enable Import Address Filter: {}",
            option_value_to_string(payload.enable_import_address_filter())
        );
        println!(
            "  Audit Enable Import Address Filter: {}",
            option_value_to_string(payload.audit_enable_import_address_filter())
        );
        println!(
            "  Enable ROP Stack Pivot: {}",
            option_value_to_string(payload.enable_rop_stack_pivot())
        );
        println!(
            "  Audit Enable ROP Stack Pivot: {}",
            option_value_to_string(payload.audit_enable_rop_stack_pivot())
        );
        println!(
            "  Enable ROP Caller Check: {}",
            option_value_to_string(payload.enable_rop_caller_check())
        );
        println!(
            "  Audit Enable ROP Caller Check: {}",
            option_value_to_string(payload.audit_enable_rop_caller_check())
        );
        println!(
            "  Enable ROP Sim Exec: {}",
            option_value_to_string(payload.enable_rop_sim_exec())
        );
        println!(
            "  Audit Enable ROP Sim Exec: {}",
            option_value_to_string(payload.audit_enable_rop_sim_exec())
        );
        println!(
            "  Override Enable Export Address Filter: {}",
            payload.override_enable_export_address_filter()
        );
        println!(
            "  Override Enable Export Address Filter Plus: {}",
            payload.override_enable_export_address_filter_plus()
        );
        println!(
            "  Override Enable Import Address Filter: {}",
            payload.override_enable_import_address_filter()
        );
        println!(
            "  Override Enable ROP Stack Pivot: {}",
            payload.override_enable_rop_stack_pivot()
        );
        println!(
            "  Override Enable ROP Caller Check: {}",
            payload.override_enable_rop_caller_check()
        );
        println!(
            "  Override Enable ROP Sim Exec: {}",
            payload.override_enable_rop_sim_exec()
        );
        if !payload.eaf_modules.is_empty() {
            println!("  EAF Modules: {:?}", payload.eaf_modules);
        }
        println!("  Flags: {}", payload.flags());
    }

    // Child Process Policy
    if let Some(child_process) = mitigations.child_process() {
        println!("\n--- Child Process Policy ---");
        println!(
            "  Disallow Child Process Creation: {}",
            option_value_to_string(child_process.disallow_child_process_creation())
        );
        println!("  Audit: {}", option_value_to_string(child_process.audit()));
        println!(
            "  Override Child Process: {}",
            child_process.override_child_process()
        );
        println!("  Flags: {}", child_process.flags());
    }

    // User Shadow Stack Policy
    if let Some(user_shadow_stack) = mitigations.user_shadow_stack() {
        println!("\n--- User Shadow Stack Policy ---");
        println!(
            "  User Shadow Stack: {}",
            option_value_to_string(user_shadow_stack.user_shadow_stack())
        );
        println!(
            "  User Shadow Stack Strict Mode: {}",
            option_value_to_string(user_shadow_stack.user_shadow_stack_strict_mode())
        );
        println!(
            "  Audit User Shadow Stack: {}",
            option_value_to_string(user_shadow_stack.audit_user_shadow_stack())
        );
        println!(
            "  Set Context IP Validation: {}",
            option_value_to_string(user_shadow_stack.set_context_ip_validation())
        );
        println!(
            "  Audit Set Context IP Validation: {}",
            option_value_to_string(user_shadow_stack.audit_set_context_ip_validation())
        );
        println!(
            "  Block Non CET Binaries: {}",
            option_value_to_string(user_shadow_stack.block_non_cet_binaries())
        );
        println!(
            "  Block Non CET Binaries Non Ehcont: {}",
            option_value_to_string(user_shadow_stack.block_non_cet_binaries_non_ehcont())
        );
        println!(
            "  Audit Block Non CET Binaries: {}",
            option_value_to_string(user_shadow_stack.audit_block_non_cet_binaries())
        );
        println!(
            "  Override User Shadow Stack: {}",
            user_shadow_stack.override_user_shadow_stack()
        );
        println!("  Flags: {}", user_shadow_stack.flags());
    }

    // SEHOP Policy
    if let Some(sehop) = mitigations.sehop() {
        println!("\n--- SEHOP Policy ---");
        println!("  Enable: {}", option_value_to_string(sehop.enable()));
        println!(
            "  Telemetry Only: {}",
            option_value_to_string(sehop.telemetry_only())
        );
        println!("  Audit: {}", option_value_to_string(sehop.audit()));
        println!("  Override SEHOP: {}", sehop.override_sehop());
        println!("  Flags: {}", sehop.flags());
    }

    // Heap Policy
    if let Some(heap) = mitigations.heap() {
        println!("\n--- Heap Policy ---");
        println!(
            "  Terminate On Error: {}",
            option_value_to_string(heap.terminate_on_error())
        );
        println!("  Override Heap: {}", heap.override_heap());
        println!("  Flags: {}", heap.flags());
    }

    println!("========================================\n");
}

fn print_help() {
    println!("Exploit Mitigations Console Tool");
    println!("Usage: program.exe <command> [arguments...]");
    println!("\nCommands:");
    println!("  help                                          - Show this help message");
    println!("  get-running <process_name>                    - Get policies from running process");
    println!("  get-registry <process_name>                   - Get policies from registry");
    println!(
        "  get-registry-json <process_name>              - Get policies from registry as JSON"
    );
    println!("  get-system                                    - Get system policy");
    println!("  get-system-json                               - Get system policy as JSON");
    println!("  get-all-processes                             - Get all process policies");
    println!("  add-process <process_name> <disable> <enable> <eaf> [force] [remove] [reset]");
    println!(
        "                                                - Add mitigations for process using numbers"
    );
    println!("  add-system <disable> <enable> <eaf> [force] [remove] [reset]");
    println!(
        "                                                - Add mitigations to system using numbers"
    );
    println!("  import-xml <file_path>                        - Import mitigations from XML");
    println!("  validate-xml <file_path>                      - Validate XML mitigations file");
    println!("  export-xml <file_path>                        - Export mitigations to XML");
    println!("\nParameters:");
    println!("  <process_name>    - Name of the process (e.g., notepad.exe)");
    println!(
        "  <disable>         - Comma-separated numbers of mitigations to disable (e.g., 0,2,15)"
    );
    println!(
        "  <enable>          - Comma-separated numbers of mitigations to enable (e.g., 1,3,16)"
    );
    println!(
        "  <eaf>             - Comma-separated EAF module names (e.g., module1.dll,module2.dll)"
    );
    println!("  [force]           - Optional: on/off/notset (default: notset)");
    println!("  [remove]          - Optional: true/false (default: false)");
    println!("  [reset]           - Optional: true/false (default: false)");
    println!("\nMitigation Numbers:");
    for (i, option) in MitigationOptions::get_all_mitigation_options()
        .iter()
        .enumerate()
    {
        println!("  {:2} - {}", i, option.to_parameter_name());
    }
}

fn parse_bool_arg(arg: &str) -> bool {
    arg.to_lowercase() == "true"
}

fn main() {
    let args: Vec<String> = env::args().collect();

    if args.len() < 2 {
        print_help();
        return;
    }

    let command = &args[1];

    match command.as_str() {
        "help" => {
            print_help();
        }
        "get-running" => {
            if args.len() < 3 {
                println!("Error: Missing process name");
                println!("Usage: {} get-running <process_name>", args[0]);
                return;
            }
            let process_name = &args[2];
            let result = Program::get_from_running_process_name(process_name);
            result.match_result(
                |processes| {
                    println!(
                        "Found {} running {} processes",
                        processes.len(),
                        process_name
                    );
                    for (i, process) in processes.iter().enumerate() {
                        println!("\n--- Running Process {} ---", i + 1);
                        print_process_mitigations(process);
                    }
                },
                |error| println!("Error: {}", error),
            );
        }
        "get-registry" => {
            if args.len() < 3 {
                println!("Error: Missing process name");
                println!("Usage: {} get-registry <process_name>", args[0]);
                return;
            }
            let process_name = &args[2];
            let result = Program::get_policy_from_registry_by_name(process_name);
            result.match_result(
                |policies| {
                    println!(
                        "Found {} registry policies for {}",
                        policies.len(),
                        process_name
                    );
                    for (i, policy) in policies.iter().enumerate() {
                        println!("\n--- Registry Policy {} ---", i + 1);
                        print_app_mitigations(policy);
                    }
                },
                |error| println!("Error: {}", error),
            );
        }
        "get-registry-json" => {
            if args.len() < 3 {
                println!("Error: Missing process name");
                println!("Usage: {} get-registry-json <process_name>", args[0]);
                return;
            }
            let process_name = &args[2];
            let result = Program::get_policy_from_registry_by_name_ex(process_name);
            result.match_result(
                |json| println!("{}", json),
                |error| println!("Error: {}", error),
            );
        }
        "get-system" => {
            let result = Program::get_system_policy();
            result.match_result(
                |policy| {
                    println!("System policy:");
                    print_app_mitigations(policy);
                },
                |error| println!("Error: {}", error),
            );
        }
        "get-system-json" => {
            let result = Program::get_system_policy_ex();
            result.match_result(
                |json| println!("{}", json),
                |error| println!("Error: {}", error),
            );
        }
        "get-all-processes" => {
            let result = Program::get_full_policies_for_all_processes();
            result.match_result(
                |policies| {
                    println!("Found {} total process policies", policies.len());
                    for (i, policy) in policies.iter().enumerate() {
                        println!("\n--- Process Policy {} ---", i + 1);
                        print_app_mitigations(policy);
                    }
                },
                |error| println!("Error: {}", error),
            );
        }
        "add-process" => {
            if args.len() < 6 {
                println!("Error: Missing required arguments");
                println!(
                    "Usage: {} add-process <process_name> <disable> <enable> <eaf> [force] [remove] [reset]",
                    args[0]
                );
                return;
            }
            let process_name = &args[2];
            let disable = &args[3];
            let enable = &args[4];
            let eaf = &args[5];
            let force = if args.len() > 6 && args[6] != "notset" {
                Some(args[6].as_str())
            } else {
                None
            };
            let remove = if args.len() > 7 {
                parse_bool_arg(&args[7])
            } else {
                false
            };
            let reset = if args.len() > 8 {
                parse_bool_arg(&args[8])
            } else {
                false
            };

            let result = Program::add_mitigations_for_process_ex(
                process_name,
                disable,
                enable,
                eaf,
                force,
                remove,
                reset,
            );
            result.match_result(
                || {
                    println!(
                        "Successfully added mitigations for process: {}",
                        process_name
                    )
                },
                |error| println!("Error: {}", error),
            );
        }
        "add-system" => {
            if args.len() < 5 {
                println!("Error: Missing required arguments");
                println!(
                    "Usage: {} add-system <disable> <enable> <eaf> [force] [remove] [reset]",
                    args[0]
                );
                return;
            }
            let disable = &args[2];
            let enable = &args[3];
            let eaf = &args[4];
            let force = if args.len() > 5 && args[5] != "notset" {
                Some(args[5].as_str())
            } else {
                None
            };
            let remove = if args.len() > 6 {
                parse_bool_arg(&args[6])
            } else {
                false
            };
            let reset = if args.len() > 7 {
                parse_bool_arg(&args[7])
            } else {
                false
            };

            let result =
                Program::add_mitigations_to_system_ex(disable, enable, eaf, force, remove, reset);
            result.match_result(
                || println!("Successfully added mitigations to system"),
                |error| println!("Error: {}", error),
            );
        }
        "import-xml" => {
            if args.len() < 3 {
                println!("Error: Missing file path");
                println!("Usage: {} import-xml <file_path>", args[0]);
                return;
            }
            let file_path = &args[2];
            let result = Program::import_xml_mitigations(file_path);
            result.match_result(
                || println!("Successfully imported XML mitigations from: {}", file_path),
                |error| println!("Error: {}", error),
            );
        }
        "validate-xml" => {
            if args.len() < 3 {
                println!("Error: Missing file path");
                println!("Usage: {} validate-xml <file_path>", args[0]);
                return;
            }
            let file_path = &args[2];
            let result = Program::validate_xml_mitigations(file_path);
            result.match_result(
                |is_valid| println!("XML validation result: {}", is_valid),
                |error| println!("Error: {}", error),
            );
        }
        "export-xml" => {
            if args.len() < 3 {
                println!("Error: Missing file path");
                println!("Usage: {} export-xml <file_path>", args[0]);
                return;
            }
            let file_path = &args[2];
            match GetProcessMitigationCommand::export_mitigation(file_path) {
                Ok(_) => println!("Successfully exported mitigations to: {}", file_path),
                Err(e) => println!("Error: {}", e),
            }
        }
        _ => {
            println!("Error: Unknown command '{}'", command);
            println!("Use 'help' to see available commands");
        }
    }
}
